from abc import ABC, abstractmethod
from typing import Any

from langchain.chains.llm import LLMChain
from langchain.chains.retrieval_qa.base import BaseRetrievalQA, RetrievalQA
from langchain.output_parsers import OutputFixingParser, PydanticOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_openai import OpenAI, ChatOpenAI

from backend.app.common.log import log
from backend.app.database.db_pinecone import rdb_conn
from backend.app.schemas.device import TenderVectorSchema


class BasePrompt(ABC):
    @abstractmethod
    def get_template(self) -> PromptTemplate:
        raise NotImplementedError()

    @abstractmethod
    def format(self, **kwargs) -> str:
        raise NotImplementedError()

    @abstractmethod
    def parser(self, result: str) -> Any:
        raise NotImplementedError()


class VectorPrompt(BasePrompt):
    """RAG增强搜索"""
    _template = """{format_instructions}
    
您将收到一条关于设备配置信息。您的目标是匹配能达到设备配置要求的设备。
您需要具备自然语言处理、数学和物理学领域的知识。
匹配过程中您需要将设备配置拆分为多个条件，然后将拆分条件逐条比对。
如果没有符合要求的设备，请返回`空`。
设备配置将包含在(```)中。
```{tender_name}; {attribute}```
符合要求的设备信息："""

    @staticmethod
    def get_out_put_parser():
        return PydanticOutputParser(pydantic_object=TenderVectorSchema)

    @staticmethod
    def format_res(result: str) -> str:
        result = result.strip()
        if result.startswith("{") and not result.endswith("}"):
            result = result + '"}'
        return result

    def get_template(self):
        parser = self.get_out_put_parser()
        format_instructions = parser.get_format_instructions()
        return PromptTemplate(
            template=self._template,
            input_variables=["attribute", "tender_name"],
            partial_variables={"format_instructions": format_instructions}
        )

    def format(self, attribute: str, tender_name: str, country: str, **kwargs) -> str:
        template = self.get_template()
        return template.format(attribute=attribute, tender_name=tender_name)

    def parser(self, result: str) -> TenderVectorSchema:
        result = self.format_res(result)

        parser = self.get_out_put_parser()
        try:
            # 使用OutputFixingParser创建一个新的解析器，该解析器能够纠正格式不正确的输出
            # new_parser = OutputFixingParser.from_llm(parser=parser, llm=ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo"))
            new_parser = OutputFixingParser.from_llm(
                parser=parser,
                llm=OpenAI(temperature=0, model_name="gpt-3.5-turbo-instruct")
            )
            # 使用新的解析器解析不正确的输出
            val = new_parser.parse(result)
            if val and val.g_id:
                val.g_id = val.g_id.strip().strip(".")
        except Exception as e:
            log.warning(f"[DevicePrompt] error: {e}")
            val = TenderVectorSchema()
        return val


class SummarizePrompt(BasePrompt):
    """对查询条件或设备属性进行摘要"""
    _template = """您是一个音视频设备专家。您将收到一个设备的属性。您的目标是对属性进行摘要。您需要将属性中符号描述转为文字描述。
设备配置将在(```)中。
```{device_meter}```

您的答案："""

    def get_template(self) -> PromptTemplate:
        return PromptTemplate(template=self._template, input_variables=["device_meter"], output_key="text")

    def format(self, text: str, **kwargs):
        tp = self.get_template()
        return tp.format(device_meter=text)

    def parser(self, result: str) -> str:
        return result


class BaseChain(ABC):
    _rdb = None
    _llm = None

    def __init__(self, namespace: str = "", index_name: str = "avl"):
        self._rdb = rdb_conn.get_rds(namespace=namespace, index_name=index_name)
        self._llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo")

    @abstractmethod
    def get_chain(self, **kwargs):
        raise NotImplementedError()

    @abstractmethod
    async def arun(self, query: str, **kwargs):
        raise NotImplementedError()


class VectorChain(BaseChain):
    """
    向量查询链
    """
    _score: float = 0.7
    _top_k: int = 3
    _chain_type: str = "map_reduce"
    _search_type: str = "similarity_score_threshold"

    def get_chain(self, shop_id: int, **kwargs) -> BaseRetrievalQA:
        rdb = self._rdb
        conditions = dict(score_threshold=self._score, k=self._top_k)
        # 国别筛选
        if shop_id:
            conditions["filter"] = dict(sid=shop_id)
        retriever_from_llm = rdb.as_retriever(
            search_type=self._search_type,
            search_kwargs=conditions
        )
        chain = RetrievalQA.from_chain_type(
            self._llm,
            chain_type=self._chain_type,
            retriever=retriever_from_llm,
            return_source_documents=True,
            verbose=True,
            chain_type_kwargs=dict(token_max=4100),
        )
        return chain

    async def arun(self, query: str, **kwargs):
        shop_id = kwargs["shop_id"]
        del kwargs["shop_id"]
        chain = self.get_chain(shop_id=shop_id, **kwargs)
        return await chain.ainvoke(input=dict(query=query))


class SummarizeChain(BaseChain):
    """总结摘要链"""

    @staticmethod
    def get_llm():
        return ChatOpenAI(temperature=0)

    def get_chain(self, prompt: PromptTemplate):
        llm = self.get_llm()
        return LLMChain(llm=llm, prompt=prompt, output_key="text")

    async def arun(self, query: str, **kwargs):
        chain = self.get_chain(prompt=kwargs["prompt"])
        return await chain.ainvoke(dict(device_meter=query))


class ChainClient:
    """客户端"""

    def __init__(self, prompt: BasePrompt, chain: BaseChain):
        self._prompt = prompt
        self._chain = chain

    async def execute(self, **kwargs):
        prompt = self._prompt
        template = prompt.get_template()
        query = prompt.format(**kwargs)
        return await self._chain.arun(prompt=template, query=query, **kwargs)

    def parser(self, result: str):
        prompt = self._prompt
        return prompt.parser(result)


async def summarize_text(text: str) -> str:
    """
    摘要：设备属性描述
    """
    prompt = SummarizePrompt()
    chain = SummarizeChain()
    client = ChainClient(prompt=prompt, chain=chain)
    result = await client.execute(text=text)
    return client.parser(result.get("text")) if result else ""


def get_vector_client() -> ChainClient:
    """获取RAG增强搜索客户端"""
    prompt = VectorPrompt()
    chain = VectorChain()
    return ChainClient(prompt=prompt, chain=chain)
